#include "njm_model.hpp"
#include <cassert>
#include <cstring>
#include <iostream>
#include <unordered_map>
#include "njm_utils.hpp"

#define TINYOBJLOADER_IMPLEMENTATION
// #include "tiny_obj_loader.h"
#include "thirdParty/tinyobjloader/tiny_obj_loader.h"

#define GLM_ENABLE_EXPERIMENTAL
#include <glm/gtx/hash.hpp>

namespace std {
    template<> struct hash<njm::NjmModel::Vertex> {
        size_t operator()(njm::NjmModel::Vertex const& vertex) const {
            size_t seed = 0;
            njm::hashCombine(seed, vertex.position, vertex.color, vertex.normal, vertex.uv);
            return seed;
        }
    };
}
namespace njm {
    

    void NjmModel::createVertexBuffers(const std::vector<Vertex> &vertices)
    {
        vCnt = static_cast<uint32_t>(vertices.size());
        assert(vCnt >= 3 && "vertex cnt must at least 3");
        VkDeviceSize bufferSize = sizeof(vertices[0]) * vCnt;
        uint32_t vertexSize = sizeof(vertices[0]);

        NjmBuffer stageBuf{
            njmDevice,
            vertexSize,
            vCnt,
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
        };
        stageBuf.map();
        stageBuf.writeToBuffer((void *)vertices.data());

        vBuf = std::make_unique<NjmBuffer>(
            njmDevice,
            vertexSize,
            vCnt,
            VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        njmDevice.copyBuffer(stageBuf.getBuffer(), vBuf->getBuffer(), bufferSize);
    }

    void NjmModel::createIndexBuffers(const std::vector<uint32_t> &indices)
    {
        iCnt = static_cast<uint32_t>(indices.size());
        hasIndexBuffer = iCnt > 0;
        if (!hasIndexBuffer) return;
        VkDeviceSize bufferSize = sizeof(indices[0]) * iCnt;
        uint32_t indexSize = sizeof(indices[0]);

        NjmBuffer stageBuf{
            njmDevice,
            indexSize,
            iCnt,
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
        };
        stageBuf.map();
        stageBuf.writeToBuffer((void *)indices.data());

        iBuf = std::make_unique<NjmBuffer>(
            njmDevice,
            indexSize,
            iCnt,
            VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        njmDevice.copyBuffer(stageBuf.getBuffer(), iBuf->getBuffer(), bufferSize);
    }

    NjmModel::NjmModel(NjmDevice& device, const NjmModel::Builder &builder): njmDevice{device}
    {
        createVertexBuffers(builder.vertices);
        createIndexBuffers(builder.indices);
    }
    NjmModel::~NjmModel()
    {
    }

    std::unique_ptr<NjmModel> NjmModel::createModelFromFile(NjmDevice& device, const std::string& filepath)
    {
        Builder builder{};
        builder.loadModel(filepath);
        std::cout << "vertex count: " << builder.vertices.size() << std::endl;
        std::cout << "index count: " << builder.indices.size() << std::endl;
        return std::make_unique<NjmModel>(device, builder);
    }

    void NjmModel::Builder::loadModel(const std::string& filepath)
    {
        tinyobj::attrib_t attrib;
        std::vector<tinyobj::shape_t> shapes;
        std::vector<tinyobj::material_t> materials;
        std::string warn, err;

        if (!tinyobj::LoadObj(&attrib, &shapes, &materials, &warn, &err, filepath.c_str())) {
            throw std::runtime_error(warn + err);
        }
        vertices.clear();
        indices.clear();

        std::unordered_map<Vertex, uint32_t> uniqueVertices{};
        for (const auto& shape : shapes) {
            for (const auto& index : shape.mesh.indices) {
                Vertex vertex{};

                if (index.vertex_index >= 0) {
                    vertex.position = {
                        attrib.vertices[3 * index.vertex_index + 0],
                        attrib.vertices[3 * index.vertex_index + 1],
                        attrib.vertices[3 * index.vertex_index + 2]
                    };
                    vertex.color = {
                        attrib.colors[3 * index.vertex_index + 0],
                        attrib.colors[3 * index.vertex_index + 1],
                        attrib.colors[3 * index.vertex_index + 2]
                    };
                }
                if (index.normal_index >= 0) {
                    vertex.normal = {
                        attrib.normals[3 * index.normal_index + 0],
                        attrib.normals[3 * index.normal_index + 1],
                        attrib.normals[3 * index.normal_index + 2]
                    };
                }
                if (index.texcoord_index >= 0) {
                    vertex.uv = {
                        attrib.texcoords[2 * index.texcoord_index + 0],
                        attrib.texcoords[2 * index.texcoord_index + 1]
                    };
                }

                if (uniqueVertices.count(vertex) == 0) {
                    uniqueVertices[vertex] = static_cast<uint32_t>(vertices.size());
                    vertices.push_back(vertex);
                }
                indices.push_back(uniqueVertices[vertex]);
            }
        }
    }

    // bind and draw
    void NjmModel::bind(VkCommandBuffer commandBuffer)
    {
        VkBuffer bufs[] = {vBuf->getBuffer()};
        VkDeviceSize offsets[] = {0};
        vkCmdBindVertexBuffers(commandBuffer, 0, 1, bufs, offsets);
        if (hasIndexBuffer) {
            vkCmdBindIndexBuffer(commandBuffer, iBuf->getBuffer(), 0, VK_INDEX_TYPE_UINT32);
        }
    }
    void NjmModel::draw(VkCommandBuffer commandBuffer)
    {
        if (hasIndexBuffer) {
            vkCmdDrawIndexed(commandBuffer, iCnt, 1, 0, 0, 0);
        } else {
            vkCmdDraw(commandBuffer, vCnt, 1, 0, 0);
        }
    }

    std::vector<VkVertexInputBindingDescription> NjmModel::Vertex::getBindingDescriptions()
    {
        std::vector<VkVertexInputBindingDescription> bindingDescriptions(1);
        bindingDescriptions[0].binding = 0;
        bindingDescriptions[0].stride = sizeof(Vertex);
        bindingDescriptions[0].inputRate = VK_VERTEX_INPUT_RATE_VERTEX;
        return bindingDescriptions;
    }
    std::vector<VkVertexInputAttributeDescription> NjmModel::Vertex::getAttributeDescriptions()
    {
        std::vector<VkVertexInputAttributeDescription> attributeDescriptions{};
        attributeDescriptions.push_back({0, 0, VK_FORMAT_R32G32B32_SFLOAT, offsetof(Vertex, position)});
        attributeDescriptions.push_back({1, 0, VK_FORMAT_R32G32B32_SFLOAT, offsetof(Vertex, color)});
        attributeDescriptions.push_back({2, 0, VK_FORMAT_R32G32B32_SFLOAT, offsetof(Vertex, normal)});
        attributeDescriptions.push_back({3, 0, VK_FORMAT_R32G32_SFLOAT, offsetof(Vertex, uv)});
        return attributeDescriptions;
    }

}